'''
	Author：Yassin
	用于选取财务报表数据，包含函数 Sel_from_bs, Sel_from_ir, Sel_from_cdr
'''
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def Sel_from_bs(L:list, RangeL:str, RangeR:str, Stock_code:list = [], Final = 1) -> pd.DataFrame:
	'''
		从资产负债表选取数据
		L:list 为且仅为需要选取的项目 code 的列表，
		RangeL:str 和 RangeR:str 为需要选取内容的时间范围，格式为 YYYYMMDD
		Stock_code : List 为所需要选取的股票的代码，默认 = []， 表示查询所有股票
		Final:int = 0 表示查询包括季报在内的所有项目，默认 = 1 表示只查询年报，默认为 1
	'''
	df = pd.read_csv("~/Desktop/mpacc2/data_invest/temp/data_lib/financial_report/balance_report/FS_Combas.csv")
	df["Accper"] = pd.to_datetime(df["Accper"])

	## 筛选需要的项目
	L = ["Stkcd", "Accper", "Typrep"] + L
	df = df[L]

	## 筛选时间范围
	df["month"] = df["Accper"].apply(lambda x:x.month)
	RangeL = pd.to_datetime(RangeL,format = "%Y%m%d")
	RangeR = pd.to_datetime(RangeR,format = "%Y%m%d")

	df = df.query("Typrep == 'A'").reset_index(drop = True)
	#df = df.query("Accper >= RangeL and Accper <= RangeR").reset_index(drop = True)
	df = df[df["Accper"] >= RangeL]
	df = df[df["Accper"] <= RangeR]
	df = df.reset_index(drop = True)

	## 筛选是否为年报
	if(Final == 1):
		df = df.query("month == 12").reset_index(drop = True)

	## 筛选股票代码
	if len(Stock_code) != 0:
		new_df = pd.DataFrame(columns = df.columns)
		for item in Stock_code:
			temp = df[df["Stkcd"] == item]
			new_df = pd.concat([new_df, temp])
		df = new_df.reset_index(drop = True)
	print(df)
	return df


def Sel_from_ir(L:list, RangeL:str, RangeR:str, Stock_code:list = [], Final = 1) -> pd.DataFrame:
	'''
		从利润表选取数据
		L:list 为且仅为需要选取的项目 code 的列表，
		RangeL:str 和 RangeR:str 为需要选取内容的时间范围，格式为 YYYYMMDD
		Stock_code : List 为所需要选取的股票的代码，默认 = []， 表示查询所有股票
		Final:int = 0 表示查询包括季报在内的所有项目，默认 = 1 表示只查询年报，默认为 1
	'''
	df = pd.read_csv("~/Desktop/mpacc2/data_invest/temp/data_lib/financial_report/income_report/FS_Comins.csv")
	df["Accper"] = pd.to_datetime(df["Accper"])

	## 筛选需要的项目
	L = ["Stkcd", "Accper", "Typrep"] + L
	df = df[L]

	## 筛选时间范围
	df["month"] = df["Accper"].apply(lambda x:x.month)
	RangeL = pd.to_datetime(RangeL,format = "%Y%m%d")
	RangeR = pd.to_datetime(RangeR,format = "%Y%m%d")

	df = df.query("Typrep == 'A'").reset_index(drop = True)
	df = df[df["Accper"] >= RangeL]
	df = df[df["Accper"] <= RangeR]
	df = df.reset_index(drop = True)

	## 筛选是否为年报
	if(Final == 1):
		df = df.query("month == 12").reset_index(drop = True)

	## 筛选股票代码
	if len(Stock_code) != 0:
		new_df = pd.DataFrame(columns = df.columns)
		for item in Stock_code:
			temp = df[df["Stkcd"] == item]
			new_df = pd.concat([new_df, temp])
		df = new_df.reset_index(drop = True)
	print(df)
	return df



def Sel_from_cdr(L:list, RangeL:str, RangeR:str, Stock_code:list = [], Final = 1) :
	'''
		从现金流量表选取数据
		L:list 为且仅为需要选取的项目 code 的列表，
		RangeL:str 和 RangeR:str 为需要选取内容的时间范围，格式为 YYYYMMDD
		Stock_code : List 为所需要选取的股票的代码，默认 = []， 表示查询所有股票
		Final:int = 0 表示查询包括季报在内的所有项目，默认 = 1 表示只查询年报，默认为 1
	'''
	df = pd.read_csv("~/Desktop/mpacc2/data_invest/temp/data_lib/financial_report/cash_direct_reprot/FS_Comscfd.csv")
	df["Accper"] = pd.to_datetime(df["Accper"])

	## 筛选需要的项目
	L = ["Stkcd", "Accper", "Typrep"] + L
	df = df[L]

	## 筛选时间范围
	df["month"] = df["Accper"].apply(lambda x:x.month)
	RangeL = pd.to_datetime(RangeL,format = "%Y%m%d")
	RangeR = pd.to_datetime(RangeR,format = "%Y%m%d")

	df = df.query("Typrep == 'A'").reset_index(drop = True)
	#df = df.query("Accper >= RangeL and Accper <= RangeR").reset_index(drop = True)
	df = df[df["Accper"] >= RangeL]
	df = df[df["Accper"] <= RangeR]
	df = df.reset_index(drop = True)

	## 筛选是否为年报
	if(Final == 1):
		df = df.query("month == 12").reset_index(drop = True)

	## 筛选股票代码
	if len(Stock_code) != 0:
		new_df = pd.DataFrame(columns = df.columns)
		for item in Stock_code:
			temp = df[df["Stkcd"] == item]
			new_df = pd.concat([new_df, temp])
		df = new_df.reset_index(drop = True)
	print(df)
	return df




if __name__ == "__main__":
	'''
		调试用内容
	'''
	Sel_from_ir(["B001100000"],20150101,20230228, Stock_code=[600519], Final = 0)
